import cv2
import numpy as np
import time

last_output_time_sidewalk = 0
last_output_time_traffic_light = 0
last_output_time_yellow_line = 0

def detect_sidewalk(frame):
    global last_output_time_sidewalk
    h = frame.shape[0]
    y_threshold = h / 3
    frame = frame[int(y_threshold):h, :]
    gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
    _, binary = cv2.threshold(gray, 200, 255, cv2.THRESH_BINARY)
    contours, _ = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    min_pixel_count = 100
    rect_count = 0
    for contour in contours:
        perimeter = cv2.arcLength(contour, True)
        approx = cv2.approxPolyDP(contour, 0.02 * perimeter, True)
        if len(approx) == 4:
            x, y, w, h = cv2.boundingRect(approx)
            rect_area = np.count_nonzero(binary[y:y + h, x:x + w])
            if rect_area > min_pixel_count:
                rect_count += 1
                cv2.drawContours(frame, [approx], -1, (0, 255, 0), 2)
    if rect_count > 4 and time.time() - last_output_time_sidewalk >= 10:
        print("检测到人行道")
        last_output_time_sidewalk = time.time()
        return True
    return False

def detect_traffic_light(frame):
    global last_output_time_traffic_light
    h = frame.shape[0]
    y_upper_threshold = h * 1 / 3
    y_lower_threshold = h * 3 / 4
    frame = frame[int(y_upper_threshold):int(y_lower_threshold), :]
    template = cv2.imread('model2.png')
    h, w = template.shape[:2]
    result = cv2.matchTemplate(frame, template, cv2.TM_SQDIFF_NORMED)
    min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(result)
    top_left = min_loc
    bottom_right = (top_left[0] + w, top_left[1] + h)
    traffic_light_area = frame[top_left[1]:bottom_right[1], top_left[0]:bottom_right[0]]
    hsv = cv2.cvtColor(traffic_light_area, cv2.COLOR_BGR2HSV)
    lower_red1 = np.array([165, 155, 90])
    upper_red1 = np.array([180, 255, 255])
    lower_green = np.array([70, 160, 50])
    upper_green = np.array([85, 255, 230])
    lower_yellow = np.array([5, 125, 120])
    upper_yellow = np.array([20, 255, 255])

    red_mask = cv2.inRange(hsv, lower_red1, upper_red1)
    red_pixel_count = np.count_nonzero(red_mask)
    green_mask = cv2.inRange(hsv, lower_green, upper_green)
    green_pixel_count = np.count_nonzero(green_mask)
    yellow_mask = cv2.inRange(hsv, lower_yellow, upper_yellow)
    yellow_pixel_count = np.count_nonzero(yellow_mask)

    if red_pixel_count > 10 and red_pixel_count < 100:
        ratio_green_yellow = green_pixel_count / yellow_pixel_count if yellow_pixel_count != 0 else 0

        if 0.1 <= ratio_green_yellow <= 1.0 and time.time() - last_output_time_traffic_light >= 10:
            print("检测到红绿灯")
            last_output_time_traffic_light = time.time()
            return True
    cv2.rectangle(frame, top_left, bottom_right, (0, 0, 255), 2)
    return False

def detect_yellow_line(frame):
    global count,last_output_time_yellow_line
    h = frame.shape[0]
    y_upper_threshold = h * 1 / 2
    y_lower_threshold = h * 3 / 4

    frame_yellow_line = frame[int(y_upper_threshold):int(y_lower_threshold), :]
    kernel = np.ones((5, 5), np.uint8)
    frame_yellow_line = cv2.dilate(frame_yellow_line, kernel, iterations=2)
    hsv_image = cv2.cvtColor(frame_yellow_line, cv2.COLOR_BGR2HSV)
    lower_yellow = np.array([20, 0, 215])
    upper_yellow = np.array([80, 100, 255])
    mask = cv2.inRange(hsv_image, lower_yellow, upper_yellow)
    yellow_pixel_count = np.sum(mask == 255)
    if yellow_pixel_count > 3000 and yellow_pixel_count < 12000:
        contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        for contour in contours:
            perimeter = cv2.arcLength(contour, True)
            approx = cv2.approxPolyDP(contour, 0.02 * perimeter, True)
            if 3 < len(approx) < 5:
                x, y, w, h = cv2.boundingRect(approx)
                aspect_ratio = w / h
                if 2 < aspect_ratio and w > 60 and h < 60:
                    cv2.rectangle(frame, (x, y + int(y_upper_threshold)), (x + w, y + h + int(y_upper_threshold)), (255, 0, 0), 2)
                    if time.time() - last_output_time_yellow_line >= 10:
                        print("检测到黄线")
                        last_output_time_yellow_line = time.time()
                        count += 1
                        if count == 3:
                            print("停止智能车")
                            return True
    return False

if __name__ == '__main__':
    video = cv2.VideoCapture('video.avi')
    fps = video.get(cv2.CAP_PROP_FPS)
    delay = int(1000 / fps)
    count = 0
    while video.isOpened():
        ret, frame = video.read()
        if ret is True:
            frame = cv2.flip(frame, 1)
            detect_sidewalk(frame)
            detect_traffic_light(frame)
            detect_yellow_line(frame)
            cv2.imshow('Frame', frame)
            if cv2.waitKey(delay) & 0xFF == ord('q'):
                break
        else:
            break
    video.release()
    cv2.destroyAllWindows()